
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier


def load_data(path):
    data = pd.read_csv(path)

    return data


def cleaning(data_param):
    data_param['Family'] = data_param['SibSp'] + data_param['Parch']
    data_param.drop(['PassengerId', 'Ticket', 'Cabin', 'SibSp', 'Parch'],
                    axis=1, inplace=True)

    # age
    average_age_titanic = data_param['Age'].mean()
    std_age_titanic = data_param['Age'].std()
    count_nan_age_titanic = data_param['Age'].isnull().sum()
    rand_1 = np.random.randint(average_age_titanic - std_age_titanic,
                               average_age_titanic + std_age_titanic,
                               size=count_nan_age_titanic)

    data_param['Age'][np.isnan(data_param['Age'])] = rand_1
    data_param['Age'] = data_param['Age'].astype(int)

    # train dataset
    data_param['Embarked'].fillna('S', inplace=True)

    # test dataset
    data_param['Fare'].fillna(data_param['Fare'].mean(), inplace=True)

    return data_param


def category_to_number(data_param):
    # Mapping Age
    data_param.loc[data_param['Age'] <= 8, 'Age'] = 0
    data_param.loc[(data_param['Age'] > 8) & (data_param['Age'] <= 16), 'Age'] = 1
    data_param.loc[(data_param['Age'] > 16) & (data_param['Age'] <= 24), 'Age'] = 2
    data_param.loc[(data_param['Age'] > 24) & (data_param['Age'] <= 32), 'Age'] = 3
    data_param.loc[(data_param['Age'] > 32) & (data_param['Age'] <= 40), 'Age'] = 4
    data_param.loc[(data_param['Age'] > 40) & (data_param['Age'] <= 48), 'Age'] = 5
    data_param.loc[(data_param['Age'] > 48) & (data_param['Age'] <= 56), 'Age'] = 6
    data_param.loc[(data_param['Age'] > 56) & (data_param['Age'] <= 64), 'Age'] = 7
    data_param.loc[(data_param['Age'] > 64) & (data_param['Age'] <= 72), 'Age'] = 8
    data_param.loc[data_param['Age'] > 72, 'Age'] = 9

    # Mapping Fare
    data_param.loc[data_param['Fare'] <= 7.91, 'Fare'] = 0
    data_param.loc[(data_param['Fare'] > 7.91) & (data_param['Fare'] <= 14.454), 'Fare'] = 1
    data_param.loc[(data_param['Fare'] > 14.454) & (data_param['Fare'] <= 31), 'Fare'] = 2
    data_param.loc[data_param['Fare'] > 31, 'Fare'] = 3
    data_param['Fare'] = data_param['Fare'].astype(int)

    # Mapping
    data_param['Sex'] = data_param['Sex'].map({'male': 0, 'female': 1})
    data_param['Embarked'] = data_param['Embarked'].map({'Q': 0, 'S': 1, 'C': 2})

    # Mapping name
    data_param['Title'] = data_param['Name'].str.extract(' ([A-Za-z]+)\.', expand=False)

    data_param['Title'] = data_param['Title'].replace('Mlle', 'Miss')
    data_param['Title'] = data_param['Title'].replace('Ms', 'Miss')
    data_param['Title'] = data_param['Title'].replace('Mme', 'Mrs')
    data_param['Title'] = data_param['Title'].replace(['Lady', 'Countess', 'Capt', 'Col',
                                                       'Don', 'Dr', 'Major', 'Rev', 'Sir', 'Jonkheer', 'Dona'], 'Rare')

    data_param['Title'] = data_param['Title'].map({"Mr": 1, "Miss": 2, "Mrs": 3, "Master": 4, "Rare": 0})
    data_param.drop('Name', axis=1, inplace=True)

    return data_param


if __name__ == '__main__':
    train_dp = './Titanic/train.csv'
    test_dp = './Titanic/test.csv'

    # training dataset
    X_train = load_data(train_dp)
    X_train = cleaning(X_train)
    X_train = category_to_number(X_train)

    # slipt train and validation
    # train = df.sample(frac=0.8, random_state=200)
    # test = df.drop(train.index)
    X_train, X_validation = train_test_split(X_train, test_size=0.15)
    y_train = X_train['Survived']
    X_train.drop('Survived', axis=1, inplace=True)
    y_validation = X_validation['Survived']
    X_validation.drop('Survived', axis=1, inplace=True)

    # testing dataset
    X_test = load_data(test_dp)
    pid = X_test['PassengerId']
    X_test = cleaning(X_test)
    X_test = category_to_number(X_test)

    print(X_train.head())
    print("--------------------------------")
    print(y_train.head())
    print("--------------------------------")

    ########################### train ###########################

    # model = SVC()
    model = RandomForestClassifier(n_estimators=400)
    model.fit(X_train, y_train)

    # validation
    acc_tr = model.score(X_train, y_train)
    acc = model.score(X_validation, y_validation)
    print(acc_tr)
    print(acc)

    # test
    pred = model.predict(X_test)

    prediction = pd.DataFrame(pred)
    prediction = pd.concat([pid, prediction], axis=1)
    prediction.columns = ['PassengerId', 'Survived']

    prediction.to_csv("./submission.csv", index=False)
